import base64
import random
import string
import struct
import time
from Crypto.Cipher import AES

class CryptoCore:
    """核心加密解密功能"""
    
    def __init__(self, encoding_aes_key, receive_id):
        """
        初始化加密核心
        
        Args:
            encoding_aes_key: 企业微信的EncodingAESKey
            receive_id: 接收者ID（自建应用用CorpID，企业微信客服用OpenKfId）
        """
        # 注意：企业微信要求AESKey = Base64_Decode(EncodingAESKey + "=")
        self.key = base64.b64decode(encoding_aes_key + '=')
        self.receive_id = receive_id

    def pkcs7_pad(self, text):
        """PKCS#7补位"""
        block_size = 32
        pad_amount = block_size - (len(text) % block_size)
        pad = chr(pad_amount)
        return text + pad * pad_amount

    def pkcs7_unpad(self, text):
        """去除PKCS#7补位"""
        pad = ord(text[-1])
        return text[:-pad]

    def encrypt(self, text, nonce=None, timestamp=None):
        """
        加密消息
        
        Args:
            text: 要加密的文本
            nonce: 随机字符串，如果为None则自动生成
            timestamp: 时间戳，如果为None则自动生成
            
        Returns:
            加密后的base64字符串
        """
        if not nonce:
            nonce = ''.join(random.sample(string.ascii_letters + string.digits, 16))
        if not timestamp:
            timestamp = str(int(time.time()))
        
        # 生成16字节随机字符串
        random_str = ''.join(random.sample(string.ascii_letters + string.digits, 16))
        
        # 消息长度（4字节，使用大端序）
        msg_len = struct.pack('>I', len(text))
        
        # 拼接：random(16字节) + msg_len(4字节) + msg + receive_id
        content = random_str.encode('utf-8') + msg_len + text.encode('utf-8') + self.receive_id.encode('utf-8')
        
        # PKCS#7补位
        content = self.pkcs7_pad(content.decode('latin1'))
        
        # AES加密
        iv = self.key[:16]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        encrypted = cipher.encrypt(content.encode('utf-8'))
        
        return base64.b64encode(encrypted).decode()

    def decrypt(self, encrypted):
        """
        解密消息 - 按照企业微信官方文档
        
        Args:
            encrypted: 加密的base64字符串
            
        Returns:
            解密后的文本，失败返回None
        """
        try:
            # Base64解码
            aes_msg = base64.b64decode(encrypted)
            
            # AES解密
            cipher = AES.new(self.key, AES.MODE_CBC, self.key[:16])
            rand_msg = cipher.decrypt(aes_msg)
            
            # 去除PKCS#7补位
            pad = rand_msg[-1]
            rand_msg = rand_msg[:-pad]
            
            # 去掉前16个随机字节
            content = rand_msg[16:]
            
            # 取出4字节的msg_len（注意字节序）
            msg_len_bytes = content[:4]
            msg_len = struct.unpack('>I', msg_len_bytes)[0]  # 使用大端序
            
            # 截取msg_len长度的msg
            msg = content[4:4+msg_len]
            
            # 剩余字节为receiveid
            receiveid = content[4+msg_len:]
            
            # 验证receiveid是否为receive_id
            if receiveid.decode('utf-8') != self.receive_id:
                raise ValueError(f"receiveid不匹配: {receiveid.decode('utf-8')} != {self.receive_id}")
            
            return msg.decode('utf-8')
            
        except Exception as e:
            print(f"解密失败: {str(e)}")
            return None 